import re
from metagpt.actions import Action


class NaiveTranslateCode(Action):
    #! Take as input an line by line annotated C/C++ code and return a Rust code that is a translation of the input C/C++ code. Write the translated code to a file named `main.rs`.
    
    name: str = "NaiveTranslateCode"
    
    ROLE_DESCRIPTION: str = """You are an expert in both C/C++ and Rust programs. You are required to translate the input C/C++ program into a Rust program. Think step by step, and eventually aggregate your thoughts in an improved version of the code. Do not ignore any function, class, or macro rule defined in the code or the content of them. Do not add new code that does not exist in the original C/C++ code, such as printing results. You should give a translation that obey the features of Rust language. The position where variables are defined and the order of declarations should not be adjusted. You must translate macro definitions in C code into Rust's macros or inline function in Rust code.  Response must be ended with ```rust \n your_code_here \n ```. """
    
    PROMPT_TEMPLATE: str = """{ROLE_DESCRIPTION}
    ## Input
    {input_code}
    ## Your Output
    """
    

    async def run(self, input_code: str):
        
        with open("main.cpp", "w") as f:
            f.write(input_code)
            
        prompt = self.PROMPT_TEMPLATE.format(ROLE_DESCRIPTION=self.ROLE_DESCRIPTION, input_code=input_code)

        rsp = await self._aask(prompt)

        code_text = NaiveTranslateCode.parse_code(rsp)
        
        with open("main.rs", "w") as f:
            f.write(code_text)

        return "main.cpp;main.rs"
    
    @staticmethod
    def parse_code(rsp):
        pattern = r"```rust(.*)```"
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text
    